#include "StdAfx.h"
#include "H264Decode.h"

#include "Ffmpeg.h"
#include <utility>



#undef min

#pragma warning (disable: 4996)


#include "WinTrace.h"
#include "TStringCast.h"
#include "TStringUtil.h"
#include "TDumpFile.h"



H264Decode::H264Decode():
    m_hwnd(),
    m_pContext(NULL),
    m_pCodec(NULL),
    m_pFrame(NULL),
    m_pScaleContext(NULL),
    m_wndWidth(),
    m_wndHeight(),
    m_pPicture(NULL),
    m_frameCount(),
    m_decBuffer(),
    m_decCallback(),
    m_pDecUser(NULL),
    m_drawCallback(),
    m_pDrawUser(NULL),
    m_oldProc(NULL)
{

    memset(&m_bmpInfo, 0, sizeof(m_bmpInfo));
    m_bmpInfo.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
    m_bmpInfo.bmiHeader.biCompression = BI_RGB;
    m_bmpInfo.bmiHeader.biBitCount = 32;
    m_bmpInfo.bmiHeader.biPlanes = 1;

}

H264Decode::~H264Decode()
{
}

void H264Decode::startup()
{
    av_register_all();
}

void H264Decode::cleanup()
{

}

bool H264Decode::open()
{
    m_pCodec = avcodec_find_decoder(AV_CODEC_ID_H264);
    if (!m_pCodec)
    {
        return false;
    }

    m_pContext = avcodec_alloc_context3(m_pCodec);
    if (avcodec_open2(m_pContext, m_pCodec, NULL) < 0)
    {
        av_free(m_pContext);
        m_pContext = NULL;
        m_pCodec = NULL;
    }

    m_pFrame = av_frame_alloc();

    m_pPicture = new AVPicture();
    memset(m_pPicture, 0, sizeof(AVPicture));

    return true;
}

void H264Decode::close()
{
    if (!m_pContext)
    {
        return;
    }

    cancelSubclass();

    av_free(m_pFrame);
    m_pFrame = NULL;

    avpicture_free(m_pPicture);
    delete m_pPicture;
    m_pPicture = NULL;

    closeScaleContext();

    avcodec_close(m_pContext);

    av_free(m_pContext);
    m_pContext = NULL;

    m_pCodec = NULL;

    m_frameCount = 0;
}

bool H264Decode::isOpen() const
{
    return (m_pCodec != NULL);
}

bool H264Decode::setVideoWnd(HWND hwnd)
{
    m_hwnd = hwnd;

    cancelSubclass();
    subclass();

    return true;
}

bool H264Decode::inputData(const unsigned char* buffer, size_t length)
{
	/*
    static size_t s_count = 0;
    if (s_count < 100)
    {
        s_count ++;
        char filename[MAX_PATH];
        sprintf(filename, "H:\\my\\%02d.h264", s_count);
        DumpFile::write((const char*)buffer, length, filename, false);
    }
    */

    AVPacket packet;
    av_init_packet(&packet);

    packet.data = (unsigned char*)buffer;
    packet.size = length;

    int got_picture = 0;
    int ret = avcodec_decode_video2(m_pContext, m_pFrame, &got_picture, &packet);
    if (got_picture)
    {
        doDecCallback();

        updateResolution();

        drawFrame();

        m_frameCount ++;
    }
    return (ret >= 0);
}

bool H264Decode::snap(const char* filename)
{
    HANDLE hFile = CreateFile(filename, GENERIC_WRITE, 0,
            NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
    if (hFile == INVALID_HANDLE_VALUE)
    {
        WinTrace("can not open file: %s\n", filename);
        return false;
    }

    int width = m_pContext->width;
    int height = m_pContext->height;
    
    AVPicture picture;
    memset(&picture, 0, sizeof(picture));
    avpicture_alloc(&picture, AV_PIX_FMT_RGB32, width, height);

    SwsContext* pScaleContext = sws_getContext(width, height,
        m_pContext->pix_fmt,
        width, height, AV_PIX_FMT_RGB32, SWS_BILINEAR, 0, 0, 0);

    sws_scale(pScaleContext, m_pFrame->data, m_pFrame->linesize, 0, height,
        picture.data, picture.linesize);

    sws_freeContext(pScaleContext);

    BITMAPINFO bmpInfo;
    memset(&bmpInfo, 0, sizeof(bmpInfo));
    bmpInfo.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
    bmpInfo.bmiHeader.biWidth = width;
    bmpInfo.bmiHeader.biHeight = height;
    bmpInfo.bmiHeader.biCompression = BI_RGB;
    bmpInfo.bmiHeader.biBitCount = 32;
    //bmpInfo.bmiHeader.biSizeImage = bmpInfo.bmiHeader.biWidth * bmpInfo.bmiHeader.biHeight * bmpInfo.bmiHeader.biBitCount /8;
    bmpInfo.bmiHeader.biPlanes = 1;
    bmpInfo.bmiHeader.biXPelsPerMeter = 0;
    bmpInfo.bmiHeader.biYPelsPerMeter = 0;
    bmpInfo.bmiHeader.biClrImportant = 0;
    bmpInfo.bmiHeader.biClrUsed = 0;

    int nColors = 0;
    BITMAPFILEHEADER hdr;
    hdr.bfType = ((WORD) ('M'<<8) |'B'); //is always "RM"
    hdr.bfSize = bmpInfo.bmiHeader.biSize + bmpInfo.bmiHeader.biSizeImage + sizeof(hdr);
    hdr.bfReserved1 = 0;		
    hdr.bfReserved2 = 0;
    hdr.bfOffBits=(DWORD) (sizeof(hdr) + bmpInfo.bmiHeader.biSize + nColors * sizeof(RGBQUAD));
    
    DWORD written = 0;
    WriteFile(hFile, &hdr, sizeof(hdr), &written, NULL);
    WinTrace("write hdr. written: %d\n", written);
    
    //fwrite(&bmpInfo.bmiHeader, 1, bmpInfo.bmiHeader.biSize, pFile);
    BOOL ret = WriteFile(hFile, &bmpInfo.bmiHeader, sizeof(bmpInfo.bmiHeader), &written, NULL);
    WinTrace("write bmpInfo. written: %d\n", written);
    if (ret == FALSE)
    {
        WinTrace("failed to write file. error:%d\n", GetLastError());
    }

    int nBytesPerLine = bmpInfo.bmiHeader.biBitCount/8*bmpInfo.bmiHeader.biWidth;
    char *pOffset = (char*)picture.data[0] + (bmpInfo.bmiHeader.biHeight-1)*nBytesPerLine;
    for (int i=0; i<bmpInfo.bmiHeader.biHeight; i++) 
    {
        WriteFile(hFile, pOffset, nBytesPerLine, &written, NULL);
        //fwrite(pOffset, 1, nBytesPerLine, pFile);
        pOffset -= nBytesPerLine;
    }

    ::CloseHandle(hFile);

    avpicture_free(&picture);

    return true;
}

bool H264Decode::getPictureSize(int& width, int& height)
{
    if (!m_pContext)
    {
        return false;
    }

    width = m_pContext->width;
    height = m_pContext->height;

    return true;
}

void H264Decode::updateResolution()
{
    int width = m_pContext->width;
    int height = m_pContext->height;

    if (::IsWindow(m_hwnd))
    {
        RECT rc;
        ::GetClientRect(m_hwnd, &rc);

        if (!IsRectEmpty(&rc))
        {
            width = rc.right - rc.left;
            height = rc.bottom - rc.top;
        }
    }

    if ((width != m_wndWidth) || (height != m_wndHeight))
    {
        closeScaleContext();

        m_wndWidth = width;
        m_wndHeight = height;

        m_pScaleContext = sws_getContext(m_pContext->width, m_pContext->height,
            m_pContext->pix_fmt,
            m_wndWidth, m_wndHeight, AV_PIX_FMT_RGB32, SWS_BILINEAR, 0, 0, 0);


        if (m_pPicture)
        {
            avpicture_free(m_pPicture);
        }
        
        avpicture_alloc(m_pPicture, AV_PIX_FMT_RGB32, m_wndWidth, m_wndHeight);
    }
}

void H264Decode::closeScaleContext()
{
    if (m_pScaleContext)
    {
        sws_freeContext(m_pScaleContext);
        m_pScaleContext = NULL;
    }
}

void H264Decode::refresh()
{
    drawFrame();
}

void H264Decode::drawFrame()
{
    if (!m_pScaleContext)
    {
        return;
    }

    int ret = sws_scale(m_pScaleContext, m_pFrame->data, m_pFrame->linesize,
        0, m_pContext->height,
        m_pPicture->data, m_pPicture->linesize);

    m_bmpInfo.bmiHeader.biWidth = m_wndWidth;
    m_bmpInfo.bmiHeader.biHeight = - m_wndHeight;

    if (!::IsWindow(m_hwnd))
    {
        return;
    }

    HDC hdc = ::GetDC(m_hwnd);

    SetDIBitsToDevice(hdc, 0, 0, m_wndWidth, m_wndHeight,
        0, 0, 0, m_wndHeight,
        m_pPicture->data[0], &m_bmpInfo, DIB_RGB_COLORS);

    doDrawCallback(hdc);

    ::ReleaseDC(m_hwnd, hdc);
}

int H264Decode::getPlayedFrame()
{
    return m_frameCount;
}

void H264Decode::setDecodeCallback(H264VideoFrameCallback cb, void* pUser)
{
    InterlockedExchangePointer((void**)&m_decCallback, cb);
    InterlockedExchangePointer(&m_pDecUser, pUser);
}

void H264Decode::setDrawCallback(H264VideoDrawCallback cb, void* pUser)
{
    InterlockedExchangePointer((void**)&m_drawCallback, cb);
    InterlockedExchangePointer(&m_pDrawUser, pUser);
}

bool H264Decode::hasDecodeCallback()
{
    return (m_decCallback != NULL);
}
 
void H264Decode::doDecCallback()
{
    if (!m_decCallback)
    {
        return;
    }

    int width = m_pFrame->linesize[0];
    int height = m_pFrame->height;

    int imageSize = width * height * 3 / 2;
    m_decBuffer.ensure(imageSize * 2);

    int planSize = width * height;
    BYTE* pBuffer = m_decBuffer.getWritePtr();
    memcpy(pBuffer, m_pFrame->data[0], planSize);
    memcpy(pBuffer + planSize, m_pFrame->data[1], planSize/4);
    memcpy(pBuffer + planSize * 5 / 4, m_pFrame->data[2], planSize/4);

    (*m_decCallback)(pBuffer, imageSize,
        width, height, FRAME_FORMAT_YV12,
        m_pDecUser);
}

void H264Decode::doDrawCallback(HDC hdc)
{
    if (!m_drawCallback)
    {
        return;
    }
    
    (*m_drawCallback)(hdc, m_pDrawUser);
}


LRESULT APIENTRY SubclassWinProc(
                                  HWND hwnd, 
                                  UINT uMsg, 
                                  WPARAM wParam, 
                                  LPARAM lParam) 
{ 
    HANDLE handle = GetProp(hwnd, TEXT("H264Decode"));
    H264Decode* pDecode = (H264Decode*)handle;
    if (!pDecode)
    {
        return 0;
    }

    return pDecode->winProc(hwnd, uMsg, wParam, lParam);
} 

bool H264Decode::subclass()
{
    if (!::IsWindow(m_hwnd))
    {
        return false;
    }

    if (!m_oldProc)
    {
        SetProp(m_hwnd, TEXT("H264Decode"), this);
        m_oldProc = (WNDPROC)SetWindowLong(m_hwnd, GWL_WNDPROC, (LONG)SubclassWinProc);
    }
    return true;
}

void H264Decode::cancelSubclass()
{
    if (!m_oldProc)
    {
        return;
    }

    if (!::IsWindow(m_hwnd))
    {
        return;
    }

    SetWindowLong(m_hwnd, GWL_WNDPROC, 
        (LONG)m_oldProc);
    m_oldProc = NULL;
}

LRESULT H264Decode::winProc(HWND hwnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
{
    if (uMsg != WM_PAINT)
    {
        return CallWindowProc(m_oldProc, hwnd, uMsg, wParam, lParam); ;
    }

    onPaint(hwnd, uMsg, wParam, lParam);

    return 0;
}

void H264Decode::onPaint(HWND hwnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
{
    if (m_frameCount <= 0)
    {
        CallWindowProc(m_oldProc, hwnd, uMsg, wParam, lParam);
    }
    else
    {
        PAINTSTRUCT ps; 
        HDC hdc; 
        hdc = BeginPaint(hwnd, &ps); 
        
        EndPaint(hwnd, &ps);

        updateResolution();

        drawFrame();
    }
}
